"""
分析任务管理服务
负责分析任务的创建、执行、监控和管理
"""
import asyncio
import uuid
from typing import Dict, Any, List, Optional, AsyncGenerator
from datetime import datetime
from enum import Enum
import logging
from sqlalchemy.orm import Session
from sqlalchemy import desc

from models.analysis_task import AnalysisTask, TaskStatus, TaskType, TaskService
from models.data_source import DataSource
from models.user import User
from core.database import get_db
from core.logging_config import get_logger
from core.exceptions import TaskError, ValidationError
from services.ai_service_factory import ai_service_factory, ServiceType, create_service_request, servers_dict
from services.data_source_manager import data_source_manager
from services.data_save_service import data_save_service

logger = get_logger("task_manager")


class TaskPriority(str, Enum):
    """任务优先级"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    URGENT = "urgent"


class TaskManager:
    """任务管理器"""
    
    def __init__(self):
        self.running_tasks: Dict[int, asyncio.Task] = {}
        self.task_queue: asyncio.Queue = asyncio.Queue()
        self.max_concurrent_tasks = 3
        self._running = False
        
    async def start(self):
        """启动任务管理器"""
        if self._running:
            return
            
        self._running = True
        logger.info("任务管理器启动")
        
        # 启动任务处理器
        for i in range(self.max_concurrent_tasks):
            logger.info(f"启动任务处理器: worker-{i}")
            asyncio.create_task(self._task_worker(f"worker-{i}"))
        
        logger.info("任务管理器启动完成")
    
    async def stop(self):
        """停止任务管理器"""
        self._running = False
        
        # 等待所有运行中的任务完成
        for task_id, task in self.running_tasks.items():
            if not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        
        logger.info("任务管理器停止")
    
    async def create_task(self, task_data: Dict[str, Any], creator_id: int, db: Session) -> AnalysisTask:
        """创建分析任务"""
        try:
            # 验证数据源
            data_source_id = task_data.get("data_source_id")
            if data_source_id:
                data_source = db.query(DataSource).filter(
                    DataSource.id == data_source_id,
                    DataSource.is_active == True
                ).first()
                if not data_source:
                    raise ValidationError(message="数据源不存在或已禁用")
            
            # 创建任务
            task = AnalysisTask(
                name=task_data["name"],
                task_type=task_data["task_type"],
                execution_type=task_data.get("execution_type", "single_analysis"),
                config=task_data.get("config", {}),
                services=task_data.get("services", []),
                data_source_id=data_source_id,
                analysis_start_date=task_data.get("analysis_start_date"),
                analysis_end_date=task_data.get("analysis_end_date"),
                creator_id=creator_id,
                max_retries=task_data.get("max_retries", 3)
            )
            
            db.add(task)
            db.commit()
            db.refresh(task)
            
            logger.info(f"创建分析任务成功: {task.name} (ID: {task.id})")
            return task
            
        except Exception as e:
            logger.error(f"创建分析任务失败: {str(e)}")
            raise TaskError(message=f"创建任务失败: {str(e)}")
    
    async def execute_task(self, task_id: int, db: Session) -> bool:
        """执行分析任务"""
        try:
            # 获取任务
            task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
            if not task:
                raise TaskError(message="任务不存在")
            
            if task.status != TaskStatus.PENDING:
                raise TaskError(message=f"任务状态错误: {task.status}")
            
            # 更新任务状态
            task.start_execution()
            db.commit()
            
            # 创建异步任务
            async_task = asyncio.create_task(self._run_task(task_id, db))
            self.running_tasks[task_id] = async_task
            
            logger.info(f"任务开始执行: {task.name} (ID: {task_id})")
            return True
            
        except Exception as e:
            logger.error(f"执行任务失败: {str(e)}")
            raise TaskError(message=f"执行任务失败: {str(e)}")

    def _update_task_services(self, task):
        """
        更新参数
        :param task:
        :return:
        """
        try:
            type_list = task.task_type.split(",")
            for type_name in type_list:
                service = servers_dict.get(type_name)
                task.services.append(service)

        except Exception as e:
            logger.error(f"服务更新异常：{e}")

    async def _run_task(self, task_id: int, db: Session):
        """运行任务的具体实现"""
        try:
            # 获取任务信息
            task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
            if not task:
                logger.error(f"任务不存在: {task_id}")
                return
            
            # 获取数据源
            data = []
            if task.data_source_id:
                data =  data_source_manager.get_data(db, task.data_source_id)
                if type(data) == dict:
                    data = data.get("data", [])
            
            task.total_count = len(data)
            start_time = task.start_time
            db.commit()
            
            # 简化处理：根据任务类型执行相应的分析
            success_count = 0
            failed_count = 0

            # if len(data) > 0:
            #     # 模拟分析过程
            #     success_count = len(data)
            #     logger.info(f"任务 {task.name} 分析完成，处理数据: {success_count} 条")
            # 更新task中的services
            self._update_task_services(task)

            for service_config in task.services:
                try:
                    service_result = await self._execute_service(
                        task_id, service_config, data, db
                    )

                    if service_result.get("success", False):
                        success_count += service_result.get("success_count", 0)
                        failed_count += service_result.get("failed_count", 0)
                    else:
                        failed_count += len(data)
                except Exception as e:
                    logger.error(f"任务执行异常: {task_id}, 服务：{service_config.get('name')} 错误: {str(e)}")
            
            # 更新任务状态
            task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
            task.complete_execution(start_time, success_count, failed_count)
            db.commit()
            
            logger.info(f"任务执行完成: {task.name} (ID: {task_id})")
            
        except Exception as e:
            logger.error(f"任务执行异常: {task_id}, 错误: {str(e)}")
            
            # 更新任务状态为失败
            task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
            if task:
                task.fail_execution(str(e))
                db.commit()
        
        finally:
            # 清理运行中的任务
            if task_id in self.running_tasks:
                del self.running_tasks[task_id]
    
    async def _execute_service(self, task_id: int, service_config: Dict[str, Any], 
                             data: List[Dict], db: Session) -> Dict[str, Any]:
        """执行单个分析服务"""
        try:
            service_name = service_config.get("name")
            service_type = service_config.get("type")
            service_params = service_config.get("parameters", {})
            
            # 创建服务记录
            service_log = TaskService(
                task_id=task_id,
                service_name=service_name,
                service_type=service_type,
                input_count=len(data)
            )
            db.add(service_log)
            db.commit()
            
            # 开始执行服务
            service_log.start_service()
            start_time = service_log.start_time
            db.commit()
            
            success_count = 0
            failed_count = 0
            
            # 批量处理数据
            batch_size = service_params.get("batch_size", 10)
            for i in range(0, len(data), batch_size):
                batch_data = data[i:i + batch_size]

                # 处理批量数据
                batch_results = await self._process_batch_data(
                    service_type, batch_data, service_params
                )

                # 保存数据到数据库
                save_result = data_save_service.save_batch_results(
                    db, task_id, batch_results
                )

                # 更新统计信息
                success_count += save_result.get("success_count", 0)
                failed_count += save_result.get("failed_count", 0)

                # 更新任务进度
                progress = ((i + len(batch_data)) / len(data)) * 100
                await self._update_task_progress(task_id, progress, db)
            
            # 完成服务
            service_log.complete_service(
                start_time,
                input_count=len(data),
                output_count=success_count + failed_count,
                success_count=success_count,
                failed_count=failed_count
            )
            db.commit()
            
            return {
                "success": True,
                "success_count": success_count,
                "failed_count": failed_count
            }
            
        except Exception as e:
            logger.error(f"服务执行失败: {service_name}, 错误: {str(e)}")
            
            # 更新服务状态
            service_log.fail_service(str(e))
            db.commit()
            
            return {
                "success": False,
                "error": str(e)
            }
    
    async def _process_batch_data(self, service_type: str, batch_data: List[Dict], 
                                 params: Dict) -> Dict[str, Any]:
        """处理批量数据"""
        try:
            # 转换服务类型
            ai_service_type = None
            if service_type == "information_extraction":
                ai_service_type = ServiceType.INFORMATION_EXTRACTION
            elif service_type == "pollution_type_recognition":
                ai_service_type = ServiceType.POLLUTION_TYPE_RECOGNITION
            elif service_type == "sentiment_analysis":
                ai_service_type = ServiceType.SENTIMENT_ANALYSIS
            elif service_type == "compliance_check":
                ai_service_type = ServiceType.COMPLIANCE_CHECK
            elif service_type == "extract_info":
                ai_service_type = ServiceType.EXTRACT_INFO
            elif service_type == "duplicate_detection":
                ai_service_type = ServiceType.DUPLICATE_DETECTION
            elif service_type == "location_analysis":
                ai_service_type = ServiceType.LOCATION_ANALYSIS
            else:
                raise ValueError(f"不支持的服务类型: {service_type}")
            
            success_count = 0
            failed_count = 0
            
            # 并发处理批量数据
            tasks = []
            for item in batch_data:
                request = create_service_request(
                    ai_service_type,
                    item,
                    params
                )
                tasks.append(ai_service_factory.process_request(request))
            
            # 执行并发处理
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 统计结果
            for result in results:
                if isinstance(result, Exception):
                    failed_count += 1
                else:
                    if result.success:
                        success_count += 1
                    else:
                        failed_count += 1
            
            return {
                "success_count": success_count,
                "failed_count": failed_count,
                "results": results
            }
            
        except Exception as e:
            logger.error(f"批量数据处理失败: {str(e)}")
            return {
                "success_count": 0,
                "failed_count": len(batch_data),
                "results": []

            }
    
    async def _update_task_progress(self, task_id: int, progress: float, db: Session):
        """更新任务进度"""
        try:
            task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
            if task:
                task.update_progress(progress)
                db.commit()
        except Exception as e:
            logger.error(f"更新任务进度失败: {str(e)}")
    
    async def stop_task(self, task_id: int, db: Session) -> bool:
        """停止任务"""
        try:
            task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
            if not task:
                raise TaskError(message="任务不存在")
            
            if task.status not in [TaskStatus.RUNNING, TaskStatus.PENDING]:
                raise TaskError(message=f"任务无法停止: {task.status}")
            
            # 取消异步任务
            if task_id in self.running_tasks:
                self.running_tasks[task_id].cancel()
                del self.running_tasks[task_id]
            
            # 更新任务状态
            task.cancel_execution()
            db.commit()
            
            logger.info(f"任务已停止: {task.name} (ID: {task_id})")
            return True
            
        except Exception as e:
            logger.error(f"停止任务失败: {str(e)}")
            raise TaskError(message=f"停止任务失败: {str(e)}")
    
    async def retry_task(self, task_id: int, db: Session) -> bool:
        """重试任务"""
        try:
            task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
            if not task:
                raise TaskError(message="任务不存在")
            
            if not task.can_retry():
                raise TaskError(message="任务无法重试")
            
            # 重试任务
            task.retry()
            db.commit()
            
            # 重新执行任务
            await self.execute_task(task_id, db)
            
            logger.info(f"任务重试成功: {task.name} (ID: {task_id})")
            return True
            
        except Exception as e:
            logger.error(f"重试任务失败: {str(e)}")
            raise TaskError(message=f"重试任务失败: {str(e)}")
    
    def get_task_status(self, task_id: int, db: Session) -> Optional[Dict[str, Any]]:
        """获取任务状态"""
        try:
            task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
            if not task:
                return None
            
            return {
                "id": task.id,
                "name": task.name,
                "status": task.status,
                "progress": task.progress,
                "total_count": task.total_count,
                "success_count": task.success_count,
                "failed_count": task.failed_count,
                "start_time": task.start_time.isoformat() if task.start_time else None,
                "end_time": task.end_time.isoformat() if task.end_time else None,
                "duration": task.duration,
                "error_message": task.error_message
            }
            
        except Exception as e:
            logger.error(f"获取任务状态失败: {str(e)}")
            return None
    
    def list_tasks(self, db: Session, filters: Dict[str, Any] = None, 
                   page: int = 1, page_size: int = 20) -> Dict[str, Any]:
        """获取任务列表"""
        try:
            query = db.query(AnalysisTask)
            
            # 应用过滤条件
            if filters:
                if "status" in filters:
                    query = query.filter(AnalysisTask.status == filters["status"])
                if "task_type" in filters:
                    query = query.filter(AnalysisTask.task_type == filters["task_type"])
                if "creator_id" in filters:
                    query = query.filter(AnalysisTask.creator_id == filters["creator_id"])
                if "start_date" in filters:
                    query = query.filter(AnalysisTask.created_at >= filters["start_date"])
                if "end_date" in filters:
                    query = query.filter(AnalysisTask.created_at <= filters["end_date"])
            
            # 获取总数
            total = query.count()
            
            # 分页
            tasks = query.order_by(desc(AnalysisTask.created_at)).offset(
                (page - 1) * page_size
            ).limit(page_size).all()
            
            return {
                "total": total,
                "page": page,
                "page_size": page_size,
                "items": [task.to_dict() for task in tasks]
            }
            
        except Exception as e:
            logger.error(f"获取任务列表失败: {str(e)}")
            raise TaskError(message=f"获取任务列表失败: {str(e)}")
    
    def get_task_detail(self, task_id: int, db: Session) -> Optional[Dict[str, Any]]:
        """获取任务详情"""
        try:
            task = db.query(AnalysisTask).filter(AnalysisTask.id == task_id).first()
            if not task:
                return None
            
            # 获取服务记录
            services = db.query(TaskService).filter(
                TaskService.task_id == task_id
            ).all()
            
            return {
                "task": task.to_dict(),
                "services": [service.to_dict() for service in services]
            }
            
        except Exception as e:
            logger.error(f"获取任务详情失败: {str(e)}")
            return None
    
    async def _task_worker(self, worker_name: str):
        """任务工作线程"""
        logger.info(f"任务工作线程启动: {worker_name}")
        
        while self._running:
            try:
                # 从队列获取任务
                task_data = await asyncio.wait_for(self.task_queue.get(), timeout=1.0)
                
                # 执行任务
                task_id = task_data.get("task_id")
                db_session = next(get_db())
                
                try:
                    await self._run_task(task_id, db_session)
                finally:
                    db_session.close()
                
            except asyncio.TimeoutError:
                continue
            except Exception as e:
                logger.error(f"任务工作线程异常: {worker_name}, 错误: {str(e)}")
        
        logger.info(f"任务工作线程停止: {worker_name}")


# 全局任务管理器实例
task_manager = TaskManager()